#!/usr/bin/env python3

import sys
sys.path.append('../../../tools')
from gmcore_plot import *
from datetime import datetime, timezone

parser = argparse.ArgumentParser(description='Plot nonhyrostatic reduced sphere mountain wave test results')
parser.add_argument('-i', '--input', help='Input data file', default='ksp15_01.h0.nc')
parser.add_argument('-t', '--time-step', type=int, default=-1, help='Time step to plot')
parser.add_argument('-o', '--output', help='Output figure')
parser.add_argument('-v', '--var', help='Variable to plot', default='w')
parser.add_argument('-l', '--loop', help='Loop the plot', action='store_true')
args = parser.parse_args()

fig = plt.figure(figsize=(8, 5))

def plot(data_path, fig):
	f = addfile(args.input)
	time = datetime.fromtimestamp(f.time.isel(time=args.time_step).values.tolist() / 1e9, timezone.utc)
	if args.var == 'w':
		levels = np.linspace(-0.11, 0.11, 23)
		colors = ListedColormap([
			[0.12, 0.16, 0.33],
			[0.11, 0.24, 0.47],
			[0.13, 0.33, 0.61],
			[0.21, 0.47, 0.72],
			[0.31, 0.63, 0.82],
			[0.45, 0.75, 0.89],
			[0.60, 0.85, 0.96],
			[0.73, 0.90, 0.96],
			[0.85, 0.95, 0.98],
			[0.94, 0.98, 0.99],
			[1.00, 1.00, 1.00],
			[1.00, 1.00, 1.00],
			[0.99, 0.97, 0.88],
			[0.99, 0.92, 0.67],
			[0.99, 0.81, 0.46],
			[1.00, 0.68, 0.28],
			[0.98, 0.53, 0.22],
			[0.95, 0.38, 0.19],
			[0.89, 0.23, 0.16],
			[0.82, 0.06, 0.16],
			[0.68, 0.06, 0.15],
			[0.53, 0.06, 0.12]
		])
	
		zlev = np.linspace(0, 10, 101) * units.km
		w = interpolate_1d(zlev, f.gz_lev.isel(time=-1, lat=90) * units('m**2 / s**2') / g, f.w_lev.isel(time=-1, lat=90))
	
		ax = fig.add_subplot(111)
		ax.set_xlim(140, 220)
		plt.colorbar(ax.contourf(f.lon, zlev, w, cmap=colors, levels=levels, extend='both'), extend='both', ticks=levels[1:-1:2])
		ax.contour(f.lon, zlev, w, levels=levels, colors='black', linewidths=0.5)
	else:
		ax = fig.add_subplot(111, projection=ccrs.PlateCarree(central_longitude=180))
		plot_contour_map(ax, f[args.var].isel(time=args.time_step, lev=5), cmap='bwr')
	
	f.close()

	ax.set_title(f'Time: {time.strftime("%H:%M")}')

plot(args.input, fig)
if args.output:
	plt.savefig(args.output, dpi=300)
elif args.loop:
	loop_plot(args.input, fig, plot)
else:
	plt.tight_layout()
	plt.show()
